"""

this is not QT app
it will solve with any image
"""
from models import get_pretrained
from models.dynamic_channel import set_uniform_channel_ratio, reset_generator
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from alfred.dl.torch.common import device

torch.set_grad_enabled(False)
direction_map = {
    'smiling': '31_Smiling',
    'young': '39_Young',
    'wavy hair': '33_Wavy_Hair',
    'gray hair': '17_Gray_Hair',
    'blonde hair': '09_Blond_Hair',
    'eyeglass': '15_Eyeglasses',
    'mustache': '22_Mustache',
}
n_style_to_change = 12


def all_rgbs_to_img(all_rgbs):
    res = all_rgbs[0].permute(1, 2, 0) + 1
    res = res * 0.5 * 255
    res = res.cpu().numpy().astype(np.uint8)
    res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
    return res


def demo():
    pretrained_type = 'generator'
    config_name = 'anycost-ffhq-config-f'

    g = get_pretrained(pretrained_type, config=config_name)
    set_uniform_channel_ratio(g, 0.5)
    g.target_res = 1024
    g.to(device)

    latent = torch.randn(1, 1, 512).to(device)
    mean_style = g.mean_style(10000)

    out, all_rgbs = g(latent, return_rgbs=True, truncation=0.5,
                      truncation_style=mean_style, randomize_noise=False)
    all_rgbs = all_rgbs[-4:][-1]
    a = all_rgbs_to_img(all_rgbs)
    cv2.imshow('aa', a)

    boundaries = get_pretrained('boundary', config_name)
    print(boundaries.keys())

    # smile
    b = boundaries[direction_map['young']].view(1, 1, -1).to(device)
    print(b)
    latent_edit = latent.clone()
    latent_edit[:, :n_style_to_change] = latent_edit[:, :n_style_to_change] + b *(50)
    out, all_rgbs = g(latent_edit, return_rgbs=True, truncation=0.5,
                      truncation_style=mean_style, randomize_noise=False)
    all_rgbs = all_rgbs[-4:][-1]
    b = all_rgbs_to_img(all_rgbs)
    cv2.imshow('ada', b)
    cv2.waitKey(0)


if __name__ == '__main__':
    demo()
